import numpy as np
import heapq
from math import sin, cos, tan, pi

class HybridAStarNode:
    def __init__(self, x, y, theta, cost, parent_idx):
        self.x = x  # grid x index
        self.y = y  # grid y index
        self.theta = theta  # vehicle orientation (rad)
        self.cost = cost  # total cost
        self.parent_idx = parent_idx  # parent node index
        self.curvature = 0.0  # path curvature

    def __lt__(self, other):
        return self.cost < other.cost

class HybridAStarPlanner:
    def __init__(self, grid_map, resolution=1.0, robot_radius=1.0, min_turning_radius=6.0):
        self.grid_map = grid_map
        self.resolution = resolution
        self.robot_radius = robot_radius
        self.min_turning_radius = min_turning_radius
        self.wheelbase = 2.5  # vehicle wheelbase (m)
        self.theta_res = pi/12  # orientation resolution (15 degrees)
        self.motion = self.get_motion_model()

    def planning(self, start, goal):
        """
        Hybrid A* path planning with vehicle motion constraints
        :param start: [x, y, theta] in world coordinates
        :param goal: [x, y] in world coordinates
        :return: path coordinates and orientations [[x,y,theta], ...]
        """
        # Convert world coordinates to grid indices
        start_idx = self.world2grid(start[:2])
        goal_idx = self.world2grid(goal)
        start_theta = start[2]

        # Initialize open and closed sets
        open_set = []
        closed_set = dict()

        # Start node
        start_node = HybridAStarNode(start_idx[0], start_idx[1], start_theta, 0.0, -1)
        heapq.heappush(open_set, (start_node.cost + self.heuristic(start_idx[0], start_idx[1], goal_idx),
                                 start_node.x, start_node.y, start_node.theta, id(start_node)))
        closed_set[(start_node.x, start_node.y, self.discretize_theta(start_theta))] = start_node

        # Main loop
        while open_set:
            cost, x, y, theta, node_id = heapq.heappop(open_set)
            current = closed_set[(x, y, self.discretize_theta(theta))]

            # Check if goal is reached
            if self.is_goal_reached((x, y), goal_idx):
                return self.extract_path(closed_set, current)

            # Explore motion primitives
            for i, _ in enumerate(self.motion):
                steer_angle, step_dist = self.motion[i]
                path = self.generate_path(current, steer_angle, step_dist)

                if not path:
                    continue

                # Process each point in the generated path
                for (nx, ny, ntheta, curvature, move_cost) in path:
                    # Check if node is within grid bounds
                    if not self.is_in_grid(nx, ny):
                        continue

                    # Check if node is obstacle-free
                    if self.grid_map[ny, nx] == 1:
                        continue

                    # Calculate new cost with curvature penalty
                    new_cost = current.cost + move_cost + self.curvature_cost(curvature)
                    new_node = HybridAStarNode(nx, ny, ntheta, new_cost, (x, y, self.discretize_theta(theta)))
                    new_node.curvature = curvature

                    # Discretize orientation for closed set
                    d_ntheta = self.discretize_theta(ntheta)

                    # Add to open set if not in closed set or has lower cost
                    if (nx, ny, d_ntheta) not in closed_set or new_cost < closed_set[(nx, ny, d_ntheta)].cost:
                        heapq.heappush(open_set, (new_cost + self.heuristic(nx, ny, goal_idx), nx, ny, ntheta, id(new_node)))
                        closed_set[(nx, ny, d_ntheta)] = new_node

        return None  # No path found

    def generate_path(self, current, steer_angle, step_dist):
        """Generate path using bicycle model"""
        path = []
        x, y, theta = current.x, current.y, current.theta
        L = self.wheelbase
        max_steps = 5

        for _ in range(max_steps):
            # Calculate turning radius and curvature
            if abs(steer_angle) < 1e-6:
                curvature = 0.0
                radius = float('inf')
            else:
                radius = L / tan(steer_angle)
                curvature = 1.0 / radius

            # Check minimum turning radius constraint
            if abs(radius) < self.min_turning_radius and radius != float('inf'):
                return None

            # Update position using bicycle model
            x += step_dist * cos(theta)
            y += step_dist * sin(theta)
            theta += step_dist / L * tan(steer_angle)
            theta = self.normalize_angle(theta)

            # Convert to grid coordinates
            gx = int(round(x))
            gy = int(round(y))

            # Add to path
            path.append((gx, gy, theta, curvature, step_dist))

        return path

    def curvature_cost(self, curvature):
        """Cost function for path curvature"""
        # Higher cost for paths with higher curvature
        return abs(curvature) * 2.0

    def heuristic(self, x, y, goal):
        """Heuristic function using Euclidean distance"""
        return np.hypot(x - goal[0], y - goal[1]) * 1.0

    def is_goal_reached(self, current, goal):
        """Check if current position is within goal tolerance"""
        return np.hypot(current[0] - goal[0], current[1] - goal[1]) < 2.0

    def is_in_grid(self, x, y):
        """Check if coordinates are within grid bounds"""
        return 0 <= x < self.grid_map.shape[1] and 0 <= y < self.grid_map.shape[0]

    def world2grid(self, world_coords):
        """Convert world coordinates to grid indices"""
        return [int(world_coords[0] / self.resolution), int(world_coords[1] / self.resolution)]

    def grid2world(self, grid_indices):
        """Convert grid indices to world coordinates"""
        return [grid_indices[0] * self.resolution, grid_indices[1] * self.resolution]

    def discretize_theta(self, theta):
        """Discretize orientation for closed set"""
        return round(theta / self.theta_res) * self.theta_res

    def normalize_angle(self, angle):
        """Normalize angle to [-pi, pi]"""
        while angle > pi:
            angle -= 2 * pi
        while angle < -pi:
            angle += 2 * pi
        return angle

    def extract_path(self, closed_set, current):
        """Extract path from closed set"""
        path = []
        while current.parent_idx != -1:
            world_x, world_y = self.grid2world([current.x, current.y])
            path.append([world_x, world_y, current.theta, current.curvature])
            current = closed_set[current.parent_idx]
        world_x, world_y = self.grid2world([current.x, current.y])
        path.append([world_x, world_y, current.theta, current.curvature])
        return path[::-1]  # Reverse to get path from start to goal

    @staticmethod
    def get_motion_model():
        """Define motion primitives (steering angle, step distance)"""
        # [steering_angle (rad), step_distance (m)]
        motion = [
            [0, 1.0],           # Straight
            [pi/6, 0.5],        # Slight left turn
            [-pi/6, 0.5],       # Slight right turn
            [pi/3, 0.3],        # Sharp left turn
            [-pi/3, 0.3]        # Sharp right turn
        ]
        return motion